package ai.example.langchain4j.config;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import reactor.core.publisher.FluxSink;

import java.util.List;
import java.util.Objects;
import java.util.function.Consumer;

public class SseEmitter implements Consumer<FluxSink<String>> {

  private ChatMemory chatMemory;
  private ChatMessage userMessage;
  private StreamingChatModel streamingChatModel;

  public SseEmitter(ChatMemory chatMemory, ChatMessage userMessage, StreamingChatModel streamingChatModel) {
    this.chatMemory = chatMemory;
    this.userMessage = userMessage;
    this.streamingChatModel = streamingChatModel;
  }

  @Override
  public void accept(FluxSink<String> emitter) {
    List<ChatMessage> message;
    if (Objects.nonNull(chatMemory)) {
      chatMemory.add(userMessage);
      message = chatMemory.messages();
    } else {
      message = List.of(userMessage);
    }
    streamingChatModel.chat(message, new StreamingChatResponseHandler() {
      @Override
      public void onPartialResponse(String s) {
        emitter.next(s); // 发送
      }

      @Override
      public void onCompleteResponse(ChatResponse chatResponse) {
        emitter.complete(); // 完成流
        chatMemory.add(new AiMessage(chatResponse.aiMessage().text()));
      }

      @Override
      public void onError(Throwable error) {
        emitter.error(error);
      }
    });
  }
}

